#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/8 20:00
# @Author  : jeff
# @FileName: user_model
import sqlalchemy
from sqlalchemy.orm import declarative_base, relationship
from sqlalchemy.sql import expression

from src.model.base_model import BaseModel, db

'''
  `id` bigint unsigned NOT NULL AUTO_INCREMENT,
  `uuid` varchar(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '总任务uuid',
  `params` json NOT NULL COMMENT '总任务参数',
  `status` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT 'pending' COMMENT '总任务状态',
  `source` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '总任务来源',
  `callback_url` varchar(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '回调地址',
  `callback_type` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '回调方式',
  `callback_status` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '回调状态',
  `work_type` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '任务类型',
  `retry` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '重试次数',
  `priority` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '优先级',
  `progress` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '进度',
  `progress_type` varchar(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '进度推送方式',
  `progress_url` varchar(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT '' COMMENT '进度推送地址',
  `create_time` datetime(3) NOT NULL,
  `update_time` datetime(3) NOT NULL,
'''


class ScanWork(BaseModel):
    __tablename__ = 'port_scan_work'

    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    uuid = db.Column(db.String(128), nullable=False, unique=True)
    params = db.Column(db.JSON(), nullable=False, default=lambda: {})
    execute_time = db.Column(db.JSON())
    status = db.Column(db.String(32))
    source = db.Column(db.String(64), nullable=True)
    callback_url = db.Column(db.String(128))
    callback_type = db.Column(db.String(64))
    callback_status = db.Column(db.String(64))
    work_type = db.Column(db.String(64))
    retry = db.Column(db.Integer)
    priority = db.Column(db.Integer)
    progress = db.Column(db.Integer)
    progress_type = db.Column(db.String(191))
    progress_url = db.Column(db.String(191))
    create_time = db.Column(sqlalchemy.DateTime, server_default=sqlalchemy.func.current_timestamp())
    update_time = db.Column(sqlalchemy.DateTime, server_default=sqlalchemy.func.current_timestamp(),
                            onupdate=sqlalchemy.func.current_timestamp())

    @classmethod
    def get_pending_work(cls, status):
        try:
            db.session.expire_all()  # 让 sessio
            pending_work = cls.query.with_for_update().filter(ScanWork.status == status).order_by(
                ScanWork.priority).first()
            return pending_work
        except Exception:
            db.session.rollback()
            raise

    @classmethod
    def get_uuid_work(cls, uuid):
        try:
            db.session.expire_all()  # 让 sessio
            pending_work = cls.query.with_for_update().filter(ScanWork.uuid == uuid).order_by(
                ScanWork.priority).first()
            return pending_work
        except Exception:
            db.session.rollback()
            raise
